import pyrealsense2 as rs
import numpy as np
import cv2
import open3d as o3d
from cam_cnofig import image_size, intrinsics_matrix, CameraInfo, factor_depth

# 初始化深度相机
pipeline = rs.pipeline()
config = rs.config()

# 创建一个坐标系
coordinate_frame = o3d.geometry.TriangleMesh.create_coordinate_frame(
    size=1,  # 坐标系的大小
    origin=[0, 0, 0]  # 坐标系的原点
)


def create_point_cloud_from_depth_image(color, depth, camera, rgb_flag = False):
    assert(depth.shape[0] == camera.height and depth.shape[1] == camera.width)
    xmap = np.arange(camera.width)
    ymap = np.arange(camera.height)
    xmap, ymap = np.meshgrid(xmap, ymap)  # 创建图像网格
    points_z = depth / camera.scale       # 深度信息  使用了广播机制
   
    points_x = (xmap - camera.cx) * points_z / camera.fx  # X = (u - cx) * Z /fx
    points_y = (ymap - camera.cy) * points_z / camera.fy  # Y = (v - cx) * Y /fy
    cloud = np.stack([points_x, points_y, points_z], axis=-1)  # 沿着最后一个轴增加维度
    if rgb_flag:
        color = color[:, :, ::-1]
        cloud = np.concatenate((cloud, color), axis=2)
        # print("cloud: ", cloud.shape)
        cloud = cloud.reshape([-1, 6])
    else:
        cloud = cloud.reshape([-1, 3])
    cloud = cloud[(cloud[:, 2] > 0.2) & (cloud[:, 2] < 2)]
    return cloud




camera = CameraInfo(640, 480, intrinsics_matrix[0][0], intrinsics_matrix[1][1], intrinsics_matrix[0][2], intrinsics_matrix[1][2], factor_depth)

# 配置深度和颜色流
config.enable_stream(rs.stream.depth, 640, 480, rs.format.z16, 30)
config.enable_stream(rs.stream.color, 640, 480, rs.format.bgr8, 30)

# 开始更新配置和同步设置
pipeline.start(config)
align = rs.align(rs.stream.color)  # Create align object for depth-color alignment

try:
    while True:
        try:
        
            # 等待深度数据帧和RGB数据帧，设置等待时间为10秒
            frames = pipeline.wait_for_frames(timeout_ms=10000)
            aligned_frames = align.process(frames)
            if not aligned_frames:
                continue  # If alignment fails, go bac
            
            depth_frame = frames.get_depth_frame()
            color_frame = frames.get_color_frame()
            
            if not depth_frame or not color_frame:
                continue

            # 获取深度图像的原始数据
            depth_data = np.asanyarray(depth_frame.get_data())
 
            # 获取RGB图像的原始数据
            color_data = np.asanyarray(color_frame.get_data())
 
            cloud = create_point_cloud_from_depth_image(color_data, depth_data, camera, rgb_flag=True)
            num_point = 10000
            cloud_size = cloud.shape[0]
            if(cloud_size >= num_point):
                idxs = np.random.choice(cloud.shape[0] , num_point, replace=False)
            else:
                idxs1 = np.arange(cloud_size)  # 补点
                idxs2 = np.random.choice(cloud_size, num_point - cloud_size, replace=True)
                idxs = np.concatenate([idxs1, idxs2], axis=0)
                idxs = np.random.choice(cloud.shape[0] , num_point, replace=False)

            cloud = cloud[idxs]


            pcd = o3d.geometry.PointCloud()
            pcd.points = o3d.utility.Vector3dVector(cloud[:, :3]) #点云数据
            pcd.colors = o3d.utility.Vector3dVector(np.array(cloud[:, 3:], dtype=np.float32) / 255.0)  # 颜色数据
            
            # 点云可是化
            o3d.visualization.draw_geometries([pcd, coordinate_frame])
            
            # # 在图像中心上点上加上标记
            # center_x, center_y = 320, 240  # 图像中心像素坐标
            # cv2.drawMarker(color_data, (center_x, center_y), (0, 255, 0), cv2.MARKER_CROSS, markerSize=20, thickness=2)
            #
            # # 使用伪彩色映射器将深度数据转换为伪彩色图像
            # colorizer = rs.colorizer()
            # depth_colormap = np.asanyarray(colorizer.colorize(depth_frame).get_data())
            #
            # # 创建一个窗口，显示彩色图像和伪彩色深度图像
            # combined_image = np.hstack((color_data, depth_colormap))
            # cv2.imshow("Combined Image", combined_image)
            #
            # # 按Esc键退出循环
            # key = cv2.waitKey(30)
            # if key == 27 or key == ord('q'):
            #     break
        except RuntimeError as e:
            print(f"等待帧时发生错误: {e}")
 
finally:
    pipeline.stop()
    cv2.destroyAllWindows()